# 1. 单专家模型 (Single Expert Model)
from model.models import RLExpert
import torch.optim as optim
from game.game import SnakeGame
from model.replay_buffer import ReplayBuffer
from train import train


def main_single_expert():
    env = SnakeGame()
    single_expert_model = RLExpert()  # 只使用一个RL专家模型
    buffer = ReplayBuffer()
    optimizer = optim.AdamW(single_expert_model.parameters(), lr=1e-4, weight_decay=1e-5)

    # 训练过程
    train(single_expert_model, env, buffer, optimizer, episodes=2000)


if __name__ == "__main__":
    main_single_expert()
